package org.dromara.easyai.extensions.cuda;

import jcuda.Pointer;
import jcuda.Sizeof;
import jcuda.jcublas.JCublas;
import jcuda.jcudnn.JCudnn;
import jcuda.jcudnn.cudnnHandle;
import jcuda.jcudnn.cudnnTensorDescriptor;
import jcuda.runtime.JCuda;
import org.dromara.jcudax.JCudax;

import static jcuda.jcudnn.cudnnDataType.CUDNN_DATA_FLOAT;
import static jcuda.jcudnn.cudnnSoftmaxAlgorithm.CUDNN_SOFTMAX_ACCURATE;
import static jcuda.jcudnn.cudnnSoftmaxMode.CUDNN_SOFTMAX_MODE_INSTANCE;
import static jcuda.jcudnn.cudnnTensorFormat.CUDNN_TENSOR_NCHW;
import static org.dromara.easyai.extensions.cuda.MatrixUtil.makeFp32Array;
import static org.dromara.easyai.extensions.cuda.MatrixUtil.makeFp32UnitColumnMajorArray;

/**
 * @author: wen.y
 * @date: 2025/2/18
 */
public class CudaFp32Util {
	/**
	 * cuda矩阵*矩阵（列主序）
	 * @param h_A 矩阵A,列主序一维数组
	 * @param h_B 矩阵B，列主序一维数组
	 * @param M 矩阵A的行数
	 * @param K 矩阵A的列数，也是矩阵B的行数
	 * @param N 矩阵B的列数
	 * @return 结果矩阵，列主序一维数组
	 * @throws Exception
	 */
	public static float[] matrixMulMatrix(float[] h_A, float[] h_B, int M, int K, int N) throws Exception {
		return matrixGemm(h_A, h_B, null, 1.0f, 0.0f, M, K, N);
	}

	/**
	 * cuda矩阵+标量（列主序）
	 * @param h_A
	 * @param scalar
	 * @param M
	 * @param N
	 * @return
	 * @throws Exception
	 */
	public static float[] matrixAddScalar(float[] h_A, float scalar, int M, int N) throws Exception {
		return matrixAddMatrix(h_A, makeFp32Array(M, N, scalar), M, N);
	}

	/**
	 * cuda矩阵+矩阵
	 * @param h_A
	 * @param h_C
	 * @param M
	 * @param N
	 * @return
	 * @throws Exception
	 */
	public static float[] matrixAddMatrix(float[] h_A, float[] h_C, int M, int N) throws Exception {
		float[] h_B = makeFp32UnitColumnMajorArray(N);
		return matrixGemm(h_A, h_B, h_C, 1.0f, 1.0f, M, N, N);
	}

	/**
	 * General Matrix Multiply 矩阵广义乘法
	 * C=α⋅A⋅B+β⋅C
	 * @param h_A
	 * @param h_B
	 * @param h_C
	 * @param alpha
	 * @param beta
	 * @param M
	 * @param K
	 * @param N
	 * @return
	 */
	public static float[] matrixGemm(float[] h_A, float[] h_B, float[] h_C, float alpha, float beta, int M, int K, int N) {
		// 分配 GPU 内存
		Pointer d_A = new Pointer();
		Pointer d_B = new Pointer();
		Pointer d_C = new Pointer();
		JCuda.cudaMalloc(d_A, (long)M * K * Sizeof.FLOAT);
		JCuda.cudaMalloc(d_B, (long)K * N * Sizeof.FLOAT);
		JCuda.cudaMalloc(d_C, (long)M * N * Sizeof.FLOAT);

		// 将数据从主机传输到设备（以列主序传入）
		JCublas.cublasSetMatrix(M, K, Sizeof.FLOAT, Pointer.to(h_A), M, d_A, M);
		JCublas.cublasSetMatrix(K, N, Sizeof.FLOAT, Pointer.to(h_B), K, d_B, K);

		if (null == h_C) {
			h_C = new float[M * N];
		} else {
			JCublas.cublasSetMatrix(M, N, Sizeof.FLOAT, Pointer.to(h_C), M, d_C, M);
		}

		// 执行矩阵乘法 C = alpha * A * B + beta * C（以列主序传入）
		JCublas.cublasSgemm('n', 'n', M, N, K, alpha, d_A, M, d_B, K, beta, d_C, M);

		// 将结果从设备传回主机
		JCublas.cublasGetMatrix(M, N, Sizeof.FLOAT, d_C, M, Pointer.to(h_C), M);

		// 释放 GPU 内存
		JCuda.cudaFree(d_A);
		JCuda.cudaFree(d_B);
		JCuda.cudaFree(d_C);
		return h_C;
	}

	/**
	 * 矩阵缩放（列主序）
	 * @param h_A
	 * @param alpha
	 * @param M
	 * @param N
	 * @return
	 */
	public static float[] matrixScale(float[] h_A, float alpha, int M, int N) {
		// 分配 GPU 内存
		Pointer d_A = new Pointer();
		JCuda.cudaMalloc(d_A, (long)M * N * Sizeof.FLOAT);

		// 将数据从主机传输到设备（以列主序传入）
		JCublas.cublasSetMatrix(M, N, Sizeof.FLOAT, Pointer.to(h_A), M, d_A, M);

		// 使用 cublasSscal 进行矩阵与标量的乘法
		JCublas.cublasSscal(M * N, alpha, d_A, 1);

		// 将结果从设备复制回主机
		JCublas.cublasGetMatrix(M, N, Sizeof.FLOAT, d_A, M, Pointer.to(h_A), M);


		// 释放 GPU 内存
		JCuda.cudaFree(d_A);

		return h_A;
	}

	/**
	 * 矩阵softmax（行主序）
	 * @param h_A
	 * @param M
	 * @param N
	 * @return
	 */
	public static float[] matrixSoftmax(float[] h_A, int M, int N) {
		// 初始化 cuDNN
		cudnnHandle handle = new cudnnHandle();
		JCudnn.cudnnCreate(handle);

		// 分配 GPU 内存
		Pointer d_A = new Pointer();
		Pointer d_C = new Pointer();
		JCuda.cudaMalloc(d_A, (long)M * N * Sizeof.FLOAT);
		JCuda.cudaMalloc(d_C, (long)M * N * Sizeof.FLOAT);

		// 将数据从主机传输到设备（以列主序传入）
		JCublas.cublasSetMatrix(M, N, Sizeof.FLOAT, Pointer.to(h_A), M, d_A, M);

		// 创建张量描述符
		// 输入
		cudnnTensorDescriptor inputDesc = new cudnnTensorDescriptor();
		JCudnn.cudnnCreateTensorDescriptor(inputDesc);
		JCudnn.cudnnSetTensor4dDescriptor(
				inputDesc,
				CUDNN_TENSOR_NCHW,     // 数据格式
				CUDNN_DATA_FLOAT,      // 数据类型
				M, 1, 1, N       // N = rows, C = 1, H = 1, W = N
		);

		// 输出
		cudnnTensorDescriptor outputDesc = new cudnnTensorDescriptor();
		JCudnn.cudnnCreateTensorDescriptor(outputDesc);
		JCudnn.cudnnSetTensor4dDescriptor(
				outputDesc,
				CUDNN_TENSOR_NCHW,     // 数据格式
				CUDNN_DATA_FLOAT,      // 数据类型
				M, 1, 1, N       // N = rows, C = 1, H = 1, W = N
		);


		// 调用 cuDNN SoftMax 函数
		JCudnn.cudnnSoftmaxForward(
				handle,
				CUDNN_SOFTMAX_ACCURATE,         // 精确模式
				CUDNN_SOFTMAX_MODE_INSTANCE,    // 对每个实例（行）进行 SoftMax
				Pointer.to(new float[]{1.0f}),  // alpha
				inputDesc,
				d_A,
				Pointer.to(new float[]{0.0f}),  // beta
				outputDesc,
				d_C
		);

		// 将结果从设备复制到主机
		float[] h_C = new float[M * N];
		JCublas.cublasGetMatrix(M, N, Sizeof.FLOAT, d_C, M, Pointer.to(h_C), M);

		// 清理资源
		JCudnn.cudnnDestroyTensorDescriptor(inputDesc);
		JCuda.cudaFree(d_A);
		JCuda.cudaFree(d_C);
		JCudnn.cudnnDestroy(handle);
		return h_C;
	}

	/**
	 * 矩阵SoftMaxPd（行主序）
	 * @param qkt
	 * @param errorMatrix
	 * @param grMatrix
	 * @param x
	 * @param y
	 * @param wordVectorDimension
	 */
	public static void matrixSoftMaxPd(float[] qkt, float[] errorMatrix, float[] grMatrix, int x, int y, float wordVectorDimension) {
		JCudax.matrixSoftMaxPdFp32(qkt, errorMatrix, grMatrix, x, y, wordVectorDimension);
	}
}
